# pip install -qU "langchain[anthropic]" to call the model

from langgraph.prebuilt import create_react_agent
from langchain_core.runnables import RunnableConfig
from langgraph.prebuilt.chat_agent_executor import AgentState
from langchain_core.messages import AnyMessage
from langchain_ollama import ChatOllama
from langgraph.checkpoint.memory import InMemorySaver
from pydantic import BaseModel

checkpointer = InMemorySaver()

class WeatherResponse(BaseModel):
    content: str

llm = ChatOllama(model="qwen3:8b", temperature=0.2, reasoning=False)

def get_weather(city: str) -> str:
    """Get weather for a given city."""
    return f"It's always raining in {city}!"

def prompt(state: AgentState, config: RunnableConfig) -> list[AnyMessage]:  
    user_name = config["configurable"].get("user_name")
    system_msg = f"You are a helpful assistant. Address the user as {user_name}."
    return [{"role": "system", "content": system_msg}] + state["messages"]

agent = create_react_agent(
    # model="anthropic:claude-3-7-sonnet-latest",
    model=llm,  # Use the ChatOllama model instead of the Anthropic model
    tools=[get_weather],
    # prompt="You are a helpful assistant"
    prompt=prompt,
    # checkpointer=checkpointer,
    response_format=WeatherResponse
)

my_config = {"configurable": {"thread_id": "1", "user_name": "John Smith"}}

for chunk in agent.stream(
    {"messages": [{"role": "user", "content": "what is the weather in sf"}]},
    config=my_config,
    # stream_mode="updates"
):
    print("-----------------------------------------")
    print(chunk)
    print("\n")


# Run the agent
# response1 = agent.invoke(
#     input={"messages": [{"role": "user", "content": "what is the weather in xi'an"}]},
#     # input={"messages": [{"role": "user", "content": "who are you"}]},
#     config=my_config
# )
# print(response1['structured_response'])  # Output: It's always raining in xi'an!

# listMsgs = response1['messages']
# for msg in listMsgs:
#     print(">>>>>>>>>>>>>>>> response1 >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
#     print(msg)

# response2 = agent.invoke(
#     input={"messages": [{"role": "user", "content": "what about new york?"}]},
#     config=my_config
# )

# listMsgs = response2['messages']
# for msg in listMsgs:
#     print(">>>>>>>>>>>>>>>>> response2 >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
#     print(msg)

